#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 11:07
# @Author  : Tingkyen
# @File    : aco蚁群算法.py
# @Software: PyCharm
# 导入需要用到的工具包
import numpy as np
import random  # 让蚂蚁能随机选择路线
import matplotlib.pyplot as plt  # 最后画个图看看结果长啥样

# ----------------------------
# 第一步：设置一些参数
# ----------------------------

NUM_CITIES = 20  # 我们让蚂蚁帮忙规划去8个城市送快递
NUM_ANTS = 50  # 派10只蚂蚁去试不同的路线
MAX_GEN = 50  # 让蚂蚁们试50轮，越多越准但越慢
ALPHA = 1  # 蚂蚁有多“迷信老路”？越大越相信气味
BETA = 2  # 蚂蚁有多“怕远路”？越大越讨厌绕远
RHO = 0.3  # 气味挥发速度：30%的气味每轮会消失（防止死记老路）
Q = 100  # 奖励值：走得好就给香香奖励！


# ----------------------------
# 第二步：生成城市位置和计算距离
# ----------------------------

def generate_cities(n):
    """随机生成n个城市的坐标，比如(10,20)表示第1个城在x=10,y=20 np.random.randint(0,100,(n,2))"""
    return np.random.rand(n, 2) * 100


def print_hi(name):
    print(f"Hi, {name}")


def calculate_distance_matrix(cities):
    """算出任意两个城市之间的距离，存成一个表格"""
    distance_matrix = np.zeros((len(cities), len(cities)))
    for i in range(len(cities)):
        for j in range(len(cities)):
            distance_matrix[i][j] = np.linalg.norm(cities[i] - cities[j])

    return distance_matrix


# ----------------------------
# 第三步：蚂蚁怎么选下一站？（靠“香味”+“近”决定）
# ----------------------------

def select_next_city(current_city, unvisited, pheromone, dist_matrix, alpha, beta):
    """
    当前蚂蚁在哪个城市，还有哪些没去过，根据“气味”和“远近”决定下一步去哪
    输入：
    current_city：当前城市编号 [0,n)
    unvisited：还没去过的城市编号 [0,n)
    pheromone：信息素维度[n,n]
    dist_matrix：距离维度[n,n]
    alpha：迷信老路的程度
    beta：怕远路的程度
    """
    probabilities = []  # 存每个可去城市的“吸引力得分”
    candidates = []  # 存可以去的城市编号
    for city in unvisited:
        if dist_matrix[current_city][city] == 0:
            continue
        # 信息素
        tau = pheromone[current_city][city] ** alpha
        # 距离
        eta = 1 / dist_matrix[current_city][city] ** beta
        # 综合得分 = 香味 × 近程度
        score = tau * eta
        probabilities.append(score)
        candidates.append(city)

    # 根据概率随机选择下一个点
    total_score = sum(probabilities)
    if total_score == 0:
        """
        算法逻辑问题：如果在算法逻辑中存在错误，导致某些情况下吸引力得分计算的结果为 0，也可能导致这个问题。比如，如果在计算信息素时，某些路径上的信息素被错误地设置为 0，
        或者在更新信息素时，某些路径上的信息素被错误地设置为 0，也可能导致这个问题。
        这种情况可能发生在以下几种情景下：
            所有的候选城市都被访问过了：这意味着蚂蚁已经走过了所有的城市，但是在实际算法逻辑中，这种情况应该已经被处理了（即在调用 select_next_city 函数之前就确认还有未访问的城市），所以通常不会导致这个问题。
            
            距离矩阵（dist_matrix）存在问题：例如，如果两个城市之间的距离被错误地设置为 0（除了城市到自身的距离之外），并且此时信息素（pheromone）也恰好是 0 或者非常小，使得即使经过 ALPHA 和 BETA 加权后，结果仍然为 0。
            
            参数配置问题：如果 ALPHA 和 BETA 参数配置不当，可能会导致某些情况下吸引力得分计算的结果为 0。比如，如果 ALPHA 和 BETA 都设置得过小，以至于信息素和启发式信息（距离倒数）的影响几乎可以忽略不计，
                        理论上也可能出现这种情况，但实际上由于初始时信息素水平较高，这种情况较为罕见。
            
            初始化问题：如果信息素矩阵（pheromone）初始化不当，或者在更新过程中出现了意外情况，导致某些路径上的信息素降到了 0，同时这些路径对应的启发式信息（距离倒数）也不足以弥补这一不足，也可能造成吸引力得分为 0。
            
            为了避免程序在这种情况下崩溃，代码中包含了检查 total_score == 0 的逻辑，并在这种情形下随机选择一个候选城市作为下一个访问目标。这实际上是一种退化策略，保证了算法能够继续运行下去。不过，在正常情况下，这种极端情况不应该发生。
        """
        # 防止全0导致崩溃
        return random.choice(candidates)
    probs = [p / total_score for p in probabilities]
    # 下一个地点
    next_city = np.random.choice(candidates, p=probs)
    return next_city


# ----------------------------
# 第四步：一只蚂蚁走一圈，完成自己的路线
# ----------------------------

def construct_path(pheromone, dist_matrix, n_cities, alpha, beta):
    """让一只蚂蚁从随机起点出发，一步步走到所有城市，返回完整路径"""
    start = np.random.randint(0, n_cities - 1)  # 随机选一个城市作为起点
    path = [start]  # 记录走过的城市
    unvisited = set(range(n_cities))  # 未走过的地方
    unvisited.remove(start)

    while len(unvisited) > 0:
        current_city = path[-1]
        next_city = select_next_city(current_city, unvisited, pheromone, dist_matrix, alpha, beta)
        path.append(next_city)
        unvisited.remove(next_city)
        # current_city = next_city

    return path


# ----------------------------
# 第五步：算一条路线总共走了多远
# ----------------------------

def path_length(path, dist_matrix):
    """输入一条路线，算总路程是多少公里"""
    length = 0
    n = len(path)
    for i in range(n):
        from_city = path[i]
        # 最后一站要回到起点
        to_city = path[(i + 1) % n]
        length += dist_matrix[from_city][to_city]
    return length


# ----------------------------
# 第六步：主函数！开始模拟蚂蚁找路
# ----------------------------

def aco_tsp():
    """蚁群算法主程序"""

    # 1. 生成城市位置和距离矩阵
    cities = generate_cities(NUM_CITIES)
    dist_matrix = calculate_distance_matrix(cities)

    # 2. 初始化信息素矩阵
    pheromone = np.ones((NUM_CITIES, NUM_CITIES)) * 0.1  # 所有边初始气味都是0.1

    # 3. 记录最好的结果
    best_length = float('inf')  # 当前最短距离（刚开始是无限大）
    best_path = None  # 当前最佳路线
    history = []  # 记录每一轮的最短距离，方便画图

    # 4. 开始迭代：让蚂蚁们跑50轮
    for gen in range(MAX_GEN):
        print(f"第{gen}轮")

        # 4.1 让10只蚂蚁去试不同的路线
        all_paths = []  # 存这一轮所有蚂蚁走的路线和长度
        for i in range(NUM_ANTS):
            path = construct_path(pheromone, dist_matrix, NUM_CITIES, ALPHA, BETA)
            length = path_length(path, dist_matrix)
            # 记录最短距离和路线
            all_paths.append((path, length))
            # 如果这只蚂蚁走得更短，更新全局最好成绩
            if length < best_length:
                best_length = length
                best_path = path.copy()
                # best_path = [int(city) for city in path.copy()]  # 明确转为整数列表

        # 记录这一轮的最好成绩
        history.append(best_length)

        # 4.2 更新信息素
        pheromone *= (1 - RHO)  # 每轮都有30%的气味挥发
        for path, length in all_paths:
            path_len = len(path)
            for i in range(path_len):
                from_city = path[i]
                to_city = path[(i+1) % path_len]
                pheromone[from_city][to_city] += Q/length  # 如果只计算从from_city到to_city的结果pheromone可能是是个上三角或者下三角矩阵，
                # 如果不加会怎样？那会出现什么情况？
                # 下次有蚂蚁想从
                # b走到a，发现这条路“没味道”或“味道很淡”。
                # 它就不会倾向于走这条路，哪怕之前有人从
                # a→b 走过并且路线很好。
                # 相当于：你知道去公司的路很顺，但回程却假装不知道——这不合理！
                pheromone[to_city][from_city] += Q/length  # 来回都加（无向图）

        # 打个招呼，看看进度
        if gen % 10 == 0:
            print(f"第 {gen} 轮：目前最短路线长度 = {best_length:.1f} 公里")
    return cities, best_length, best_path, history


# ----------------------------
# 第七步：画图展示结果
# ----------------------------

def plot_result(cities, best_path, history):
    """画出最终路线和收敛过程"""
    plt.figure(figsize=(12, 5))

    # 左边：城市路线图
    plt.subplot(1, 2, 1)
    plt.title("蚂蚁找到的最佳路线")
    x = cities[:, 0]
    y = cities[:, 1]
    plt.scatter(x, y, c='blue')  # 所有城市标出来

    # 把最佳路线连成线
    route = best_path + [best_path[0]]  # 最后连回起点
    for i in range(len(route) - 1):
        from_city = cities[route[i]]
        to_city = cities[route[i + 1]]
        plt.plot([from_city[0], to_city[0]], [from_city[1], to_city[1]], 'g-')

    # 标上城市编号
    for i in range(len(cities)):
        plt.text(cities[i][0], cities[i][1], f" {i}", color="red", fontsize=12)

    # 右边：每一轮的最短距离变化
    plt.subplot(1, 2, 2)
    plt.title("路线越来越短")
    plt.plot(history, 'b-o', markersize=3)
    plt.xlabel("轮数")
    plt.ylabel("最短路线长度")
    plt.grid(True)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    print("🐜 蚂蚁军团启动！正在寻找最优快递路线...")

    # 运行蚁群算法
    cities, best_length, best_path, history = aco_tsp()

    print(f"\n🎉 成功！蚂蚁找到了一条超短路线：")
    print(f"路线: {' → '.join(map(str, best_path))} → 回到 {best_path[0]}")
    # print(f"路线: {' → '.join([str(int(city)) for city in best_path])} → 回到 {int(best_path[0])}")
    print(f"总长度: {best_length:.2f} 公里")

    # 画图看看
    plot_result(cities, best_path, history)
    print_hi("Tingkyen")
